import subprocess
import os
import io
import sys
from contextlib import redirect_stdout, redirect_stderr
from PIL import Image, ImageDraw, ImageFont

def capture_terminal_content(command, output_file="terminal_output.png", font_size=14, test_mode=False, line_height_factor=1.5):
    """
    捕获终端命令输出并保存为图片
    
    Args:
        command: 要执行的命令字符串
        output_file: 输出图片文件名
        font_size: 字体大小
        test_mode: 是否使用测试模式
        line_height_factor: 行高因子，控制行间距
    """

    # 测试模式：直接使用中文字符串而不是运行命令
    if test_mode:
        output = "这是一个中文测试文本\n中文显示测试123\nTest English"
        print("使用测试文本模式")
    else:
        # 运行命令并捕获所有输出
        process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = process.communicate()
        
        # 合并标准输出和错误输出，尝试不同的编码方式以支持中文
        encodings = ['utf-8', 'cp936', 'gb2312', 'gbk']
        output = ""
        
        # 尝试解码标准输出
        for encoding in encodings:
            try:
                output = stdout.decode(encoding)
                print(f"使用编码: {encoding}")
                break
            except UnicodeDecodeError:
                continue
        else:
            # 如果所有编码都失败，使用errors='ignore'
            output = stdout.decode('utf-8', errors='ignore')
            print("使用默认编码(可能有乱码)")
        
        # 尝试解码错误输出
        if stderr:
            error_output = ""
            for encoding in encodings:
                try:
                    error_output = stderr.decode(encoding)
                    break
                except UnicodeDecodeError:
                    continue
            else:
                error_output = stderr.decode('utf-8', errors='ignore')
            
            output += "\nERRORS:\n" + error_output
    
    # 设置字体 - 使用Windows系统字体目录
    windows_font_dir = "C:\\Windows\\Fonts\\"
    font_candidates = [
        os.path.join(windows_font_dir, "simhei.ttf"),  # 黑体
        os.path.join(windows_font_dir, "msyh.ttf"),   # 微软雅黑
        os.path.join(windows_font_dir, "simsun.ttc"),  # 宋体
    ]
    
    font = None
    for font_path in font_candidates:
        if os.path.exists(font_path):
            try:
                font = ImageFont.truetype(font_path, font_size)
                print(f"使用字体: {font_path}")
                break
            except Exception as e:
                print(f"加载字体 {font_path} 失败: {e}")
    
    # 如果所有字体都加载失败，尝试不指定字体（让系统自动处理）
    if font is None:
        print("所有字体加载失败，尝试不指定字体")
        font = None  # None表示让系统自动选择字体
    
    # 计算文本尺寸
    lines = output.splitlines()
    if not lines:
        lines = ["No output"]
    
    # 计算最大行宽和总高度
    # 使用较新的方法替代已废弃的getsize()
    try:
        # 尝试使用getbbox()方法（较新的PIL版本）
        bbox = font.getbbox("X")
        char_width = bbox[2] - bbox[0]  # 右边界 - 左边界
        char_height = bbox[3] - bbox[1]  # 下边界 - 上边界
    except AttributeError:
        # 兼容旧版本PIL
        char_width, char_height = font.getsize("X")
    
    # 设置行高因子，增加行距
    line_height = int(char_height * line_height_factor)
    
    max_line_length = max(len(line) for line in lines)
    total_lines = len(lines)
    
    # 创建图像
    padding = 50
    img_width = min(char_width * max_line_length + 2*padding, 10000)  # 限制最大宽度
    img_height = line_height * total_lines + 2*padding
    
    # 如果图像太高，分割为多个图像
    max_height = 8000  # PIL的常见限制
    if img_height > max_height:
        print(f"Output is too long, splitting into multiple images")
        num_images = (img_height // max_height) + 1
        lines_per_image = total_lines // num_images
        
        for i in range(num_images):
            start = i * lines_per_image
            end = (i+1) * lines_per_image if i < num_images-1 else total_lines
            part_lines = lines[start:end]
            
            part_img = Image.new("RGB", (img_width, line_height*len(part_lines)+2*padding), "white")
            draw = ImageDraw.Draw(part_img)
            
            y = padding
            for line in part_lines:
                draw.text((padding, y), line, font=font, fill="black")
                y += line_height
            
            part_output = output_file.replace(".png", f"_part{i+1}.png")
            part_img.save(part_output)
            print(f"Saved part {i+1} to {part_output}")
        
        return
    
    # 创建单张图像
    image = Image.new("RGB", (img_width, img_height), "white")
    draw = ImageDraw.Draw(image)
    
    y = padding
    for line in lines:
        draw.text((padding, y), line, font=font, fill="black")
        y += line_height
    
    image.save(output_file)
    print(f"Terminal output saved to {output_file}")

def capture_python_output(func, output_file="python_output.png", font_size=14, line_height_factor=1.5, *args, **kwargs):
    """
    捕获Python函数的输出并保存为图片
    
    Args:
        func: 要执行并捕获输出的函数
        output_file: 输出图片文件名
        font_size: 字体大小
        line_height_factor: 行高因子，控制行间距
        *args, **kwargs: 传递给func的参数
    """
    # 创建字符串IO对象来捕获输出
    f = io.StringIO()
    e = io.StringIO()
    
    print(f"捕获Python函数 '{func.__name__}' 的输出...")
    
    try:
        # 重定向标准输出和标准错误到字符串IO对象
        with redirect_stdout(f), redirect_stderr(e):
            # 执行函数
            func(*args, **kwargs)
        
        # 获取捕获的输出
        stdout_content = f.getvalue()
        stderr_content = e.getvalue()
        
        # 合并输出
        output = stdout_content
        if stderr_content:
            output += "\nERRORS:\n" + stderr_content
        
        print("使用Python函数输出模式")
        
        # 以下代码复用capture_terminal_content中的绘图逻辑
        # 设置字体 - 使用Windows系统字体目录
        windows_font_dir = "C:\\Windows\\Fonts\\"
        font_candidates = [
            os.path.join(windows_font_dir, "simhei.ttf"),  # 黑体
            os.path.join(windows_font_dir, "msyh.ttf"),   # 微软雅黑
            os.path.join(windows_font_dir, "simsun.ttc"),  # 宋体
        ]
        
        font = None
        for font_path in font_candidates:
            if os.path.exists(font_path):
                try:
                    font = ImageFont.truetype(font_path, font_size)
                    print(f"使用字体: {font_path}")
                    break
                except Exception as e:
                    print(f"加载字体 {font_path} 失败: {e}")
        
        # 如果所有字体都加载失败，尝试不指定字体（让系统自动处理）
        if font is None:
            print("所有字体加载失败，尝试不指定字体")
            font = None  # None表示让系统自动选择字体
        
        # 计算文本尺寸
        lines = output.splitlines()
        if not lines:
            lines = ["No output"]
        
        # 计算最大行宽和总高度
        # 使用较新的方法替代已废弃的getsize()
        try:
            # 尝试使用getbbox()方法（较新的PIL版本）
            bbox = font.getbbox("X")
            char_width = bbox[2] - bbox[0]  # 右边界 - 左边界
            char_height = bbox[3] - bbox[1]  # 下边界 - 上边界
        except AttributeError:
            # 兼容旧版本PIL
            char_width, char_height = font.getsize("X")
        
        # 设置行高因子，增加行距
        line_height = int(char_height * line_height_factor)
        
        max_line_length = max(len(line) for line in lines)
        total_lines = len(lines)
        
        # 创建图像
        padding = 50
        img_width = min(char_width * max_line_length + 2*padding, 10000)  # 限制最大宽度
        img_height = line_height * total_lines + 2*padding
        
        # 如果图像太高，分割为多个图像
        max_height = 8000  # PIL的常见限制
        if img_height > max_height:
            print(f"Output is too long, splitting into multiple images")
            num_images = (img_height // max_height) + 1
            lines_per_image = total_lines // num_images
            
            for i in range(num_images):
                start = i * lines_per_image
                end = (i+1) * lines_per_image if i < num_images-1 else total_lines
                part_lines = lines[start:end]
                
                part_img = Image.new("RGB", (img_width, line_height*len(part_lines)+2*padding), "white")
                draw = ImageDraw.Draw(part_img)
                
                y = padding
                for line in part_lines:
                    draw.text((padding, y), line, font=font, fill="black")
                    y += line_height
                
                part_output = output_file.replace(".png", f"_part{i+1}.png")
                part_img.save(part_output)
                print(f"Saved part {i+1} to {part_output}")
            
            return
        
        # 创建单张图像
        image = Image.new("RGB", (img_width, img_height), "white")
        draw = ImageDraw.Draw(image)
        
        y = padding
        for line in lines:
            draw.text((padding, y), line, font=font, fill="black")
            y += line_height
        
        image.save(output_file)
        print(f"Python函数输出已保存到 {output_file}")
        
    finally:
        # 清理
        f.close()
        e.close()

if __name__ == "__main__":
    # 示例使用 - 使用测试模式优先测试中文显示
    # 使用测试模式直接测试中文显示，使用默认行距因子1.5
    capture_terminal_content("", "test_chinese.png", test_mode=True)

    # 创建一个更大行距的版本（2.0倍行距）
    capture_terminal_content("", "test_chinese_large_spacing.png", test_mode=True, line_height_factor=2.0)

    # 然后再运行实际命令
    try:
        # 使用默认行距
        capture_terminal_content("tree", "terminal_tree.png")
        # 使用更大行距
        capture_terminal_content("tree", "terminal_tree_large_spacing.png", line_height_factor=2.0)
    except Exception as e:
        print(f"运行tree命令出错: {e}")
        # 如果tree命令失败，尝试dir命令
        capture_terminal_content("dir", "terminal_dir.png")
        capture_terminal_content("dir", "terminal_dir_large_spacing.png", line_height_factor=2.0)


    # Python输出捕获示例
    def example_python_function():
        """示例函数，用于测试Python输出捕获功能"""
        print("这是Python函数的输出")
        print("这是第二行输出，包含中文")
        print(f"当前工作目录: {os.getcwd()}")
        # 故意产生一个错误输出
        try:
            1/0
        except Exception as e:
            print(f"产生一个错误: {e}")

    # 捕获Python函数的输出并保存为图片
    print("\n测试Python函数输出捕获功能...")
    capture_python_output(example_python_function, "python_function_output.png")
    capture_python_output(example_python_function, "python_function_output_large_spacing.png", line_height_factor=2.0)